import matplotlib.pyplot as plt
import pandas as pd
import numpy as np  # 添加这行

def simulate_cleaning_data():
    topics = [
        "央视记者采访燕郊爆炸被驱离",
        "新闻学教授怒骂张雪峰",
        "胡锡进谈虐猫事件",
        "巴以冲突",
        "姜萍数学竞赛违规事件"
    ]
    return pd.DataFrame({
        '原始数据量': [12800, 9200, 10500, 8700, 7600],
        '清洗后数据量': [10400, 7800, 8600, 6900, 7200],
        '原始连贯度': [0.52, 0.58, 0.61, 0.49, 0.67],
        '清洗后连贯度': [0.78, 0.72, 0.83, 0.69, 0.81]
    }, index=topics)


def visualize_cleaning_effect(df):
    plt.rcParams.update({
        'font.family': 'SimSun',
        'axes.unicode_minus': False
    })

    fig, ax1 = plt.subplots(figsize=(12, 7), dpi=300)
    ax2 = ax1.twinx()

    x = np.arange(len(df))
    width = 0.35

    # 数据量柱状图
    bars1 = ax1.bar(x - width / 2, df['原始数据量'], width,
                    color='#3498DB', label='原始数据量')
    bars2 = ax1.bar(x + width / 2, df['清洗后数据量'], width,
                    color='#2ECC71', label='清洗后数据量')

    # 连贯度折线图
    ax2.plot(x, df['原始连贯度'], 's--', color='#E74C3C',
             markersize=10, label='原始连贯度')
    ax2.plot(x, df['清洗后连贯度'], 'D-', color='#9B59B6',
             markersize=10, label='清洗后连贯度')

    # 坐标轴设置
    ax1.set_ylabel('数据量 (条)', fontsize=12)
    ax2.set_ylabel('语义连贯度', fontsize=12)
    ax1.set_xticks(x)
    ax1.set_xticklabels(df.index, rotation=20, ha='right')
    ax1.set_ylim(0, 15000)
    ax2.set_ylim(0.4, 0.9)

    # 数据标签
    for i, topic in enumerate(df.index):
        ax2.text(i, df.loc[topic, '清洗后连贯度'] - 0.03,
                 f"+{df.loc[topic, '清洗后连贯度'] - df.loc[topic, '原始连贯度']:.2f}",
                 ha='center', color='#34495E', fontsize=9)

    # 图例
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2,
               loc='upper left', bbox_to_anchor=(0.02, 1.18),
               ncol=2, frameon=False)

    plt.title("热点事件文本清洗效果分析", fontsize=14, pad=25)
    plt.tight_layout()
    plt.savefig('hot_topics_cleaning_corrected.png')
    plt.close()


def main():
    df = simulate_cleaning_data()
    df['数据保留率'] = df['清洗后数据量'] / df['原始数据量']
    df['连贯度提升'] = df['清洗后连贯度'] - df['原始连贯度']

    print("热点事件清洗效果分析报告".center(40, '='))
    print(df[['原始数据量', '清洗后数据量', '数据保留率',
              '原始连贯度', '清洗后连贯度', '连贯度提升']].round(2))

    visualize_cleaning_effect(df)


if __name__ == "__main__":
    main()